import numpy as np
import random

def loadData(name,root_path="data/"):
    """
    :param name: 数据集名字，实验需要CORA和CITESEER两个
    :return:
    """
    names = ["CORA", "CITESEER"]
    split_ratio = 0.2
    assert name in names

    if name == "CORA":
        edgeData = root_path + "/cora/cora.cites"
        nodeData = root_path + "/cora/cora.content"
    elif name == "CITESEER":
        edgeData = root_path + "/citeseer/citeseer.cites"
        nodeData = root_path + "/citeseer/citeseer.content"

    with open(edgeData,"r") as f1, open(nodeData,"r") as f2:

        lines_edge_data = f1.readlines() # 边
        lines_node_data = f2.readlines() # 节点

        # 把节点和label转为index
        node_name_set = set()
        class_set = set()
        for line in lines_node_data:
            r = line[:-1].split("\t")
            node_name_set.add(r[0])
            class_set.add(r[-1])
        node_to_index = {name: index for index, name in enumerate(node_name_set)}
        label_to_index = {name: index for index, name in enumerate(class_set)}

        ID_labels_attr_list = [] # [节点ID，节点类别，节点特征]

        for line in lines_node_data:
            line_data = line[:-1].split("\t")
            ID_labels_attr_list.append([node_to_index[line_data[0]], label_to_index[line_data[-1]], np.array(line_data[1:-1]).astype(np.int) ] )

        ID_labels_attr_list.sort() # 这步不能少，不然node_attr和node对不上号,因为node_attr并不包含与node的映射信息，默认按升序node_index排列
        node = np.array([a[0] for a in ID_labels_attr_list]).reshape(-1, 1)
        node_attr = np.array([a[-1] for a in ID_labels_attr_list])
        ID_labels = np.array([[a[0], a[1]] for a in ID_labels_attr_list])
        np.random.shuffle(ID_labels)

        edge_list = []
        for line in lines_edge_data:
            line = line[:-1]
            l,r = line.split("\t")
            try:
                line_data = [node_to_index[l], node_to_index[r]]
            except:
                continue # 数据集似乎会缺少一些节点的特征和类别，直接跳过相关的边
            edge_list.append(line_data)
        edge = np.array(edge_list)

    trainData = {}
    testData = {}
    split_index = int(split_ratio * len(lines_node_data))

    trainData["node"] = node
    trainData["edge"] = np.vstack((edge,edge[:,[1, 0]]))
    trainData["node_attr"] = node_attr
    trainData["ID"] = ID_labels[:split_index, 0].reshape(-1, 1) # ID信息类似于mask
    trainData["label"] =ID_labels[:split_index, 1].reshape(-1, 1) # GNN输出后，按照ID取出来的结果就是对应的label

    testData["node"] = node
    testData["edge"] = np.vstack((edge,edge[:,[1, 0]]))
    testData["node_attr"] = node_attr
    testData["ID"] = ID_labels[split_index:, 0].reshape(-1, 1)

    # 数据集的一些额外信息
    testData["num_class"] = trainData["num_class"] = len(label_to_index)
    testData["num_features"] = trainData["num_features"] = node_attr.shape[1]
    testData["num_samples"] = trainData["num_samples"] = len(trainData["ID"])

    testLabel = ID_labels[split_index:, 1].reshape(-1, 1)

    return trainData, testData, testLabel